function sum_rate = calculate_sum_rate(H, s, scheme, SINR_approx, sigma2, varargin)
% CALCULATE_SUM_RATE 计算系统总速率（文献公式47）
%   输入：
%     H: [M×K]信道矩阵
%     s: [M×1]发送信号向量
%     scheme: 预编码方案 'JGP'/'PGP'/'CGP'
%     SINR_approx: 确定性SINR近似值（可选）
%     sigma2: 噪声功率
%     varargin: 方案特定参数
%   输出：
%     sum_rate: 总速率(bps/Hz)
%
%   文献依据：公式(24,32,38,47)

    % ========== 1. 参数解析 ==========
    [M, K] = size(H);
    
    % ========== 2. 方案选择 ==========
    switch scheme
        case 'JGP'
            % JGP方案（公式24）
            if nargin < 5 || isempty(SINR_approx)
                SINR = calculate_JGP_SINR(H, s, sigma2);
            else
                SINR = SINR_approx;
            end
            
        case 'PGP'
            % PGP方案（公式32）
%             groups = varargin{1};
            groups = [ones(6,1); 2*ones(6,1); 3*ones(6,1)];
            if nargin < 5 || isempty(SINR_approx)
                SINR = calculate_PGP_SINR(H, s, groups, sigma2);
            else
                SINR = SINR_approx;
            end
            
        case 'CGP'
            % CGP方案（公式38）
            common_groups = {[1,2], 3};
            groups = [ones(6,1); 2*ones(6,1); 3*ones(6,1)];
%             common_groups = varargin{1};
%             groups = varargin{2};
            if nargin < 5 || isempty(SINR_approx)
                SINR = calculate_CGP_SINR(H, s, common_groups, groups, sigma2);
            else
                SINR = SINR_approx;
            end
            
        otherwise
            error('未知预编码方案: %s', scheme);
    end
    
    % ========== 3. 计算和速率 ==========
    % 公式47：R_sum = sum_{k=1}^K log2(1 + SINR_k)
    sum_rate = sum(log2(1 + SINR));
end

%% ========== 子函数：JGP SINR计算 ==========
function SINR = calculate_JGP_SINR(H, s, sigma2)
% 计算JGP方案的瞬时SINR（公式24）
    [M, K] = size(H);
    SINR = zeros(K, 1);
    
    for k = 1:K
        % 提取用户k信道
        h_k = H(:, k);
        
        % 计算有用信号
        useful_signal = h_k' * s;
        
        % 计算干扰+噪声
        interference = 0;
        for j = 1:K
            if j == k, continue; end
            interference = interference + abs(h_k' * s)^2;
        end
        noise = sigma2;
        
        % 计算SINR
        SINR(k) = abs(useful_signal)^2 / (interference + noise);
    end
end

%% ========== 子函数：PGP SINR计算 ==========
function SINR = calculate_PGP_SINR(H, s, groups, sigma2)
% 计算PGP方案的瞬时SINR（公式32）
    [M, K] = size(H);
    G = length(unique(groups));
    SINR = zeros(K, 1);
    
    for g = 1:G
        group_users = find(groups == g);
        
        for idx = 1:length(group_users)
            k = group_users(idx);
            h_k = H(:, k);
            
            % 组内有用信号
            useful_signal = h_k' * s;
            
            % 组内干扰
            intra_interference = 0;
            for j = group_users'
                if j == k, continue; end
                intra_interference = intra_interference + abs(h_k' * s)^2;
            end
            
            % 组间干扰
            inter_interference = 0;
            for j = 1:K
                if groups(j) == g, continue; end
                inter_interference = inter_interference + abs(h_k' * s)^2;
            end
            
            % 计算SINR
            SINR(k) = abs(useful_signal)^2 / (intra_interference + inter_interference + sigma2);
        end
    end
end

%% ========== 子函数：CGP SINR计算 ==========
function SINR = calculate_CGP_SINR(H, s, common_groups, groups, sigma2)
% 计算CGP方案的瞬时SINR（公式38）
    [M, K] = size(H);
    SINR = zeros(K, 1);
    
    for gp = 1:length(common_groups)
        group_users = [];
        for g = common_groups{gp}
            group_users = [group_users; find(groups == g)];
        end
        
        for k = group_users'
            h_k = H(:, k);
            useful_signal = h_k' * s;
            
            % 公共组内干扰
            intra_interference = 0;
            for j = group_users'
                if j == k, continue; end
                intra_interference = intra_interference + abs(h_k' * s)^2;
            end
            
            % 公共组间干扰
            inter_interference = 0;
            for j = 1:K
                if ismember(j, group_users), continue; end
                inter_interference = inter_interference + abs(h_k' * s)^2;
            end
            
            % 计算SINR
            SINR(k) = abs(useful_signal)^2 / (intra_interference + inter_interference + sigma2);
        end
    end
end

%% ========== 子函数：确定性近似 ==========
function SINR_approx = get_deterministic_SINR(scheme, H_eff, P_T, alpha, varargin)
% 获取确定性SINR近似（文献公式25,33,39）
    switch scheme
        case 'JGP'
            % 公式25
            SINR_approx = JGP_deterministic_SINR(H_eff, P_T, alpha);
            
        case 'PGP'
            % 公式33
            groups = varargin{1};
            SINR_approx = PGP_deterministic_SINR(H_eff, groups, P_T, alpha);
            
        case 'CGP'
            % 公式39
            common_groups = varargin{1};
            groups = varargin{2};
            SINR_approx = CGP_deterministic_SINR(H_eff, common_groups, groups, P_T, alpha);
    end
end